#include "ap.h"

#include "xp.h"
#include "mem.h"
#include "assert.h"

#include <stdarg.h>
#include <limits.h>
#include <string.h>
#include <ctype.h>

#define T AP_T

// Use sign + absolute value to represent an sign integer.
struct T {
    int sign;
    int ndigits;
    int size;
    XP_T digits;
};

#define iszero(x) ((x)->ndigits == 1 && (x)->digits[0] == 0)
#define isone(x) ((x)->ndigits == 1 && (x)->digits[0] == 1)
#define iseven(x) (((x)->digits[0] & 1) == 0)

#define issamesign(x, y) (((x)->sign ^ (y)->sign) == 0)
#define maxdigits(x, y) ((x)->ndigits > (y)->ndigits ? (x)->ndigits : (y)->ndigits)

#define declare_and_init \
    unsigned char d[sizeof(unsigned long)]; \
    struct T t; t.size = (int)sizeof(d); t.digits = d;

static T
mk(int size)
{
    assert(size > 0);

    T z = CALLOC(1, sizeof(*z) + size);
    z->sign = 1;
    z->ndigits = 1;
    z->size = size;
    z->digits = (XP_T)(z + 1);
    return z;
}

static T
normalize(T z, int n)
{
    z->ndigits = XP_length(n, z->digits);
    return z;
}

static T
set(T z, long n)
{
    if (n == LONG_MIN)
        XP_fromint(z->size, z->digits, LONG_MAX + 1ul);
    else if (n < 0)
        XP_fromint(z->size, z->digits, -n);
    else
        XP_fromint(z->size, z->digits, n);
    z->sign = n < 0 ? -1 : 1; // for 0 is 1
    return normalize(z, z->size);
}

static int
cmp(T x, T y)
{
    if (x->ndigits != y->ndigits)
        return x->ndigits - y->ndigits;
    else
        return XP_cmp(x->ndigits, x->digits, y->digits);
}

T AP_new(long n)
{
    return set(mk((int)sizeof(long)), n);
}

void AP_free(T *z)
{
    assert(z && *z);
    FREE(*z);
}

// number = {white}[- | +]{white} digit {digit}
T AP_fromstr(const char *str, int base, char **end)
{
    assert(str);
    assert(base >= 2 && base <= 36);

    const char *p = str;
    while (*p && isspace(*p))
        p++;

    char sign = '\0';
    if (*p == '-' || *p == '+')
        sign = *p++;

    const char *start;
    int k, n = 0;
    for (; *p == '0' && p[1] == '0'; p++)
        ;
    start = p;

    // Find the number of digit in base
    for (;
         (*p >= '0' && *p <= '9' && *p < '0' + base) ||
         (*p >= 'a' && *p <= 'z' && *p < 'a' + base - 10) ||
         (*p >= 'A' && *p <= 'Z' && *p < 'A' + base - 10);
         p++) {
        n++;
    }

    // Calculate the needed bytes m
    // m = (n * lg base) / 8
    // lg base is k: 2^k >= base
    for (k = 1; (1<<k) < base; k++)
        ;
    T z = mk(((k*n + 7) & ~7) / 8);
    p = start;

    char *endp;
    int carry = XP_fromstr(z->size, z->digits, p, base, &endp);
    assert(carry == 0);
    normalize(z, z->size);

    if (endp == p) {
        endp = (char*)str;
        AP_free(&z);
        z = AP_new(0);

    } else {
        z->sign = iszero(z) || sign != '-' ? 1 : -1;
    }
    if (end)
        *end = endp;
    return z;
}

long AP_toint(T x)
{
    assert(x);

    unsigned long u = XP_toint(x->ndigits, x->digits);
    if (x->sign == -1)
        return -(long)u; // -LONG_MIN == LONG_MIN
    else
        return (long)u;
}

char* AP_tostr(char *str, int size, int base, T x)
{
    assert(x);
    assert(base >= 2 && base <= 36);
    assert(str == NULL || size > 1);

    if (str == NULL) {
        // Calculate number of characters in str
        int k;
        // 2^6 = 64 > 32(base)
        for (k = 5; (1 << k) > base; k--)
            ;
        size = (8*x->ndigits) / k + 1 + 1;
        if (x->sign == -1)
            size++;
        str = ALLOC(size);
    }

    XP_T q = ALLOC(x->ndigits);
    memcpy(q, x->digits, x->ndigits);

    if (x->sign == -1) {
        str[0] = '-';
        XP_tostr(str+1, size-1, base, x->ndigits, q);
    } else {
        XP_tostr(str, size, base, x->ndigits, q);
    }

    FREE(q);
    return str;
}

void AP_fmt(int code, va_list_box *box,
            int put(int c, void *cl), void *cl,
            unsigned char flags[256], int width, int precision)
{
    (void)code;
    assert(box && flags);

    T x = va_arg(box->ap, T);
    assert(x);

    char *buf = AP_tostr(NULL, 0, 10, x);
    Fmt_putd(buf, strlen(buf), put, cl, flags, width, precision);
    FREE(buf);
}

int AP_cmp(T x, T y)
{
    assert(x);
    assert(y);

    if (!issamesign(x, y))
        return x->sign;
    else if (x->sign == 1)
        return cmp(x, y);
    else
        return cmp(y, x);
}

int AP_cmpi(T x, long y)
{
    declare_and_init;
    return AP_cmp(x, set(&t, y));
}

T AP_neg(T x)
{
    assert(x);

    T z = mk(x->ndigits);
    memcpy(z->digits, x->digits, x->ndigits);
    z->ndigits = x->ndigits;
    z->sign = iszero(z) ? 1 : -x->sign;
    return z;
}

// Add without sign
static T
add(T z, T x, T y)
{
    int n = y->ndigits;
    if (x->ndigits < n) {
        return add(z, y, x);

    } else if (x->ndigits > n) {
        int carry = XP_add(n, z->digits, x->digits, y->digits, 0);
        z->digits[z->size-1] = XP_sum(x->ndigits-n, &z->digits[n], &x->digits[n], carry);

    } else {
        z->digits[n] = XP_add(n, z->digits, x->digits, y->digits, 0);
    }
    return normalize(z, z->size);
}

// Sub without sign
// Assert: |x| >= |y|
static T
sub(T z, T x, T y)
{
    int n = y->ndigits;
    int borrow = XP_sub(n, z->digits, x->digits, y->digits, 0);
    if (x->ndigits > n)
        borrow = XP_diff(x->ndigits-n, &z->digits[n], &x->digits[n], borrow);
    assert(borrow == 0);
    return normalize(z, z->size);
}

T AP_add(T x, T y)
{
    assert(x);
    assert(y);

    T z;
    if (issamesign(x, y)) {
        z = add(mk(maxdigits(x, y)+1), x, y);
        z->sign = iszero(z) ? 1 : x->sign;

    } else {
        if (cmp(x, y) > 0) {
            z = sub(mk(x->ndigits), x, y);
            z->sign = iszero(z) ? 1 : x->sign;
        } else {
            z = sub(mk(y->ndigits), y, x);
            z->sign = iszero(z) ? 1 : -x->sign;
        }
    }
    return z;
}

T AP_sub(T x, T y)
{
    assert(x);
    assert(y);

    T z;
    if (!issamesign(x, y)) {
        z = add(mk(maxdigits(x, y)+1), x, y);
        z->sign = iszero(z) ? 1 : x->sign;

    } else {
        if (cmp(x, y) > 0) {
            z = sub(mk(x->ndigits), x, y);
            z->sign = iszero(z) ? 1 : x->sign;
        } else {
            z = sub(mk(y->ndigits), y, x);
            z->sign = iszero(z) ? 1 : -x->sign;
        }
    }
    return z;
}

T AP_mul(T x, T y)
{
    assert(x);
    assert(y);

    T z = mk(x->ndigits + y->ndigits);
    XP_mul(z->digits, x->ndigits, x->digits, y->ndigits, y->digits);
    normalize(z, z->size);
    z->sign = iszero(z) || issamesign(x, y) ? 1: -1;
    return z;
}

static void
div(T x, T y, T *qp, T *rp)
{
    T q = mk(x->ndigits); // The max digit number
    T r = mk(y->ndigits); // The max digit number
    XP_T tmp = ALLOC(x->ndigits + y->ndigits + 2);
    XP_div(x->ndigits, q->digits, x->digits, y->ndigits, y->digits, r->digits, tmp);
    FREE(tmp);

    normalize(q, q->size);
    normalize(r, r->size);
    q->sign = iszero(q) || issamesign(x, y) ? 1 : -1;

    *qp = q;
    *rp = r;
}

// When x and y with same sign: quotient is |x|/|y|,
// When x and y with different sign: quotient is |x|/|y| or |x|/|y|+1(if |x|%|y| != 0).
T AP_div(T x, T y)
{
    assert(x);
    assert(y);
    assert(!iszero(y));

    T q, r;
    div(x, y, &q, &r);

    if (!issamesign(x, y) && !iszero(r)) {
        int carry = XP_sum(q->size, q->digits, q->digits, 1);
        assert(carry == 0);
        normalize(q, q->size);
    }

    AP_free(&r);
    return q;
}

// The remainder is always postive.
// When x and y with same sign: remainder is |x|%|y|,
// When x and y with different sign: remainder is |x|%|y| or |y|-|x|%|y|(if |x|%|y| != 0).
T AP_mod(T x, T y)
{
    assert(x);
    assert(y);
    assert(!iszero(y));

    T q, r;
    div(x, y, &q, &r);

    if (!issamesign(x, y) & !iszero(r)) {
        int borrow = XP_sub(r->size, r->digits, y->digits, r->digits, 0);
        assert(borrow == 0);
        normalize(r, r->size);
    }

    AP_free(&q);
    return r;
}

static T
mulmod(T x, T y, T p)
{
    T xy = AP_mul(x, y);
    T z = AP_mod(xy, p);
    AP_free(&xy);
    return z;
}

// Recursive computation:
//   if y is even : z = (x^(y/2)) * (x^(y/2))
//   if y not even: z = (x^(y/2)) * (x^(y/2)) * x
T AP_pow(T x, T y, T p)
{
    assert(x);
    assert(y);
    assert(y->sign == 1);
    assert(!p || (p->sign == 1 && !iszero(p) && !isone(p)));

    if (iszero(x))
        return AP_new(0);
    if (iszero(y))
        return AP_new(1);
    if (isone(x))
        return AP_new(iseven(y) ? 1 : x->sign);

    T z;
    if (p) {
        if (isone(y)) {
            z = AP_addi(x, 0);

        } else {
            // (x * y) mod p = ((x mod p) * (y mod p)) mod p
            T y2 = AP_rshift(y, 1);
            T t = AP_pow(x, y2, NULL);
            z = mulmod(t, t, p);
            AP_free(&y2);
            AP_free(&t);

            if (!iseven(y)) {
                z = mulmod(y2=AP_mod(x, p), t=z, p);
                AP_free(&y2);
                AP_free(&t);
            }
        }

    } else {
        if (isone(y)) {
            z = AP_addi(x, 0);

        } else {
            T y2 = AP_rshift(y, 1);
            T t = AP_pow(x, y2, NULL);
            z = AP_mul(t, t);
            AP_free(&y2);
            AP_free(&t);

            if (!iseven(y)) {
                z = AP_mul(x, t=z);
                AP_free(&t);
            }
        }
    }

    return z;
}

T AP_addi(T x, long y)
{
    declare_and_init;
    return AP_add(x, set(&t, y));
}

T AP_subi(T x, long y)
{
    declare_and_init;
    return AP_sub(x, set(&t, y));
}

T AP_muli(T x, long y)
{
    declare_and_init;
    return AP_mul(x, set(&t, y));
}

T AP_divi(T x, long y)
{
    declare_and_init;
    return AP_div(x, set(&t, y));
}

long AP_modi(T x, long y)
{
    declare_and_init;
    T r = AP_mod(x, set(&t, y));
    long rem = AP_toint(r);
    AP_free(&r);
    return rem;
}

T AP_lshift(T x, int s)
{
    assert(x);
    assert(s >= 0);

    T z = mk(x->ndigits + (((s+7) & ~7)) / 8);
    XP_lshift(z->size, z->digits, x->ndigits, x->digits, s, 0);
    z->sign = x->sign;
    return normalize(z, z->size);
}

T AP_rshift(T x, int s)
{
    assert(x);
    assert(s >= 0);

    if (s >= 8*x->ndigits) {
        return AP_new(0);

    } else {
        T z = mk(x->ndigits - s/8);
        XP_rshift(z->size, z->digits, x->ndigits, x->digits, s, 0);
        z->sign = iszero(z) ? 1 : x->sign;
        return normalize(z, z->size);
    }
}

